import pandas as pd
from matplotlib import pyplot as plt
from sklearn.ensemble import RandomForestClassifier
import seaborn as sns

from sklearn.impute import SimpleImputer
from utils.utils import *

class RandomForestUtil:
    def __init__(self,dp,split_random_state,pic_save_path,feature_save_path=rf_feature_path):
        self.total_df = None
        self.feature_save_path = opj(feature_save_path,str(split_random_state))
        self.pic_save_path = opj(pic_save_path,str(split_random_state),'feature_extract','rf')
        self.dp = dp


    def extract_feature(self,df,data_type,top_k=10):
        X = df.drop(['Patient_ID', target_column], axis=1)

        # 先划分训练/测试集
        X_train, X_test, y_train, y_test, _, _ = self.dp.split_train_and_test(df.copy())

        # 在训练集上 fit median 填补
        imputer = SimpleImputer(strategy='median')
        X_train_imputed = imputer.fit_transform(X_train)

        # 转换为 DataFrame
        X_train = pd.DataFrame(X_train_imputed, columns=X.columns)

        # 后续训练模型
        rf = RandomForestClassifier(n_estimators=100, random_state=42)
        rf.fit(X_train, y_train)

        # 特征重要性计算
        importance = rf.feature_importances_
        feature_names = X.columns

        feature_importance_df = pd.DataFrame({
            'Feature': feature_names,
            'Importance': importance
        }).sort_values(by='Importance', ascending=False)

        # 可视化前k个特征
        plt.figure(figsize=(10, 8))
        sns.barplot(
            data=feature_importance_df.head(top_k),
            x='Importance',
            y='Feature',
            hue='Feature',
            palette='viridis',
            dodge=False,
            legend=False
        )
        plt.title('Top {} Important Features from Random Forest'.format(top_k))
        plt.tight_layout()
        psp = opj(self.pic_save_path, data_type.replace('FeatureSummary', ''))
        md(psp)
        plt.savefig(opj(psp, 'feature_importance.pdf'))
        plt.close()

        # 返回包含前k个特征 + Patient_ID + target_column 的原始数据
        selected_features = feature_importance_df.head(top_k)['Feature'].tolist()
        df_filtered = df[['Patient_ID', target_column] + selected_features]

        md(self.feature_save_path)
        df_filtered.to_csv(opj(self.feature_save_path, data_type.replace('FeatureSummary', '') + '.csv'), index=False)
        if self.total_df is None:
            self.total_df = df_filtered
        else:
            self.total_df = self.total_df.merge(df_filtered, how='left', on=['Patient_ID', target_column])
        return False

    def save_total_df(self):
        md(self.feature_save_path)
        self.total_df.to_csv(opj(self.feature_save_path, 'total.csv'), index=False)
